# coding:utf-8
from bert_serving.client import BertClient
import subprocess
import threading
import time
import os


def exist_service():
    """
    判断当前是否已经启动了bert service
    :return:
    """
    grep_cmd = "ps -aux|grep bert-serving-start"
    _sta, output = subprocess.getstatusoutput(grep_cmd)
    if len(output) > 0:
        # exist service
        service_count = 0
        pooling_service_count = 0
        for line in output.split("\n"):
            items = line.strip().split()
            status = items[7]
            cmd = " ".join(items[10:])
            if not cmd.startswith("/home/daniel/anaconda3/envs/py3/bin/python /home/daniel/anaconda3/envs/py3/bin/bert-serving-start"):
                continue
            if "pooling_strategy" in cmd:
                start_pos = cmd.find("pooling_strategy")
                next_empty_pos = cmd.find(" ", start_pos+len("pooling_strategy")+1)
                if next_empty_pos == -1:
                    next_empty_pos = len(cmd)
                if cmd[start_pos+len("pooling_strategy")+1: next_empty_pos] == "None":
                    service_count += 1
                else:
                    pooling_service_count += 1
            else:
                pooling_service_count += 1
        return service_count, pooling_service_count
    else:
        # doesn't exist service
        return 0, 0


def start_word_service():
    """
    启动单词bert service
    :return:
    """
    while True:
        service_count, pooling_service_count = exist_service()
        if service_count > 0:
            print("service already start")
            time.sleep(60)
        else:
            start_cmd = "bert-serving-start -model_dir /home/daniel/下载/chinese_L-12_H-768_A-12 -pooling_strategy NONE " \
                "-num_worker=2"
            subprocess.getoutput(start_cmd)



def start_sentence_service():
    """
    启动句子bert service
    :return:
    """
    service_count, pooling_service_count = exist_service()
    if pooling_service_count > 0:
        "pooling service already start"
        return
    start_cmd = "bert-serving-start -model_dir /home/daniel/下载/chinese_L-12_H-768_A-12 -num_worker=2"
    t = threading.Thread(target=subprocess.getoutput, args=[start_cmd])
    t.start()


def main():
    bc = BertClient()
    doc_vecs = bc.encode(['今天天空很蓝,阳光明媚', '今天天气好晴朗', '现在天气如何', '自然语言处理', '机器学习任务'])
    print(doc_vecs)


if __name__ == "__main__":
    # pooling 和 非 pooling 无法同时启动
    # start_sentence_service()
    start_word_service()
    # s_c, ps_c = exist_service()
    # print("service_count", s_c)
    # print("pooling service count", ps_c)
